還是先AI slob 剩下的 我們之後再說
這幾天在參加活動@@
希望過幾天還有存量可以繼續QAQ
可能要有四天的 這個 先等我一下 我後續補給你ㄌ
昨天兩張表已經把位置標出來了。
跟 Dao 比,N=14 我們還輸;N=15 後來贏,而且 diff 是 0;N=16 開始付 2-pass 的 memory 錢。跟 Tensor Core 比,慢的不是 MMA,是餵資料。
今天做 Day 15 預告的最後一件事:如果分析已經說明原本寫法輸在哪裡,就把路徑重新切乾淨。這裡說的激進,不是再塞更多 trick,而是比較像 FlashAttention 那種精神——中間結果能不寫出去,就不要寫出去;kernel 邊界用 profile 反推,不靠感覺切。
N=15 其實已經在付 2-pass 的錢昨天表一的 N=15 = 82.98 us,是切完以後的數字。切之前不是這樣。
原本 N=15 會被拆成:
Pass 1:N=10 的 single-block kernel
每個 thread 只拿 1 個 float4(4 個 float)
Pass 2:剩下 5 層,走 tiled shared memory
一次只搬 TILE_COLS = 16 條 column
兩個 kernel,中間把整份 state 寫回 global memory 再讀回來。batch=128 時,這一份是:
128 × 32768 × 4 bytes = 16 MB
光這一趟 round-trip 就是多讀寫 16 MB。Pass 2 的 tile 又只有 16 條 column 寬,bandwidth 很難打滿。butterfly 層數跟切完之後一樣是 15 層;帳單幾乎全在搬資料和第二次 launch。
當時量到大約 278 us。Dao 在同一題大約 98 us。我們慢的不是加減法,是太早離開 single-block。
NCU 在 N=14 還給過一個很像結論的暗示:chunk 已經 64 KB,再往 N=15 把整份丟進 shared memory 會直接撞限制。如果只讀這句,下一步會變成「那 N=15 只好 2-pass,去優化 Pass 2」。
那就是切錯邊界。
N=15 不要拆 pass,shared memory 只拿來換位置N=15 是 32768 個 float,128 KB。它不適合全部住在 shared memory,但可以住在 256 個 thread 的 register 裡:
256 threads × 32 個 float4 = 32768 floats
每個 thread 持有 128 個 float
shared memory 大約 32 KB,只做 warp 之間的 corner turning,不當整份 state 的倉庫。這跟 Dao 的切法是同一類:butterfly 在 register / warp shuffle 裡做完,smem 只負責換 layout。
切完之後,同一合約(FP32, batch=128):
N=15 |
時間 | 相對 Dao | diff |
|---|---|---|---|
| 切之前:10+5 兩個 pass | ~278 us | 約 0.35× | 沒對齊 |
| Dao | 98.35 us | 1× | — |
| 切之後:一個 block 做完 | 82.98 us | 1.19× | 0.0e+00 |
278 us → 82.98 us,大約 3.3 倍。省掉的是那一趟 16 MB round-trip,加上一次 kernel launch,加上 Pass 2 那條太窄的 tile。diff = 0 則是因為這次連加減順序都跟 Dao 對上了。
occupancy 看起來不會漂亮。128 個 float 放 register,compiler 隨時可能 spill。但表上的 3.3 倍說明:這題裡,多一趟 global memory 比 occupancy 低更傷。要盯的是 local memory load/store 有沒有冒出來,不是看到 occupancy 低就把 kernel 拆回兩個 pass。
N=15 吃得下,Pass 2 才能變輕第一刀不只救 N=15。它改變 large-N 怎麼切。
切之前,第一段只吃到 10
N=16 → 剩下 6 層(超過 5,必須走 tiled smem)
N=20 → 剩下 10 層
切之後,第一段吃到 15
N=16 → 剩下 1 層
N=20 → 剩下 5 層(剛好一個 warp)
剩下 ≤ 5 層時,Pass 2 可以走 Day 17 那顆 shuffle kernel:每個 lane 拿一個 row,__shfl_xor_sync 做完 column butterfly。block 裡不再走 shared memory 的完整交換。
跨 block 的 global memory 仍要付。2-pass 的下限還在。但 block 內能省的,就不要再付一次。
昨天表一的 N=16 = 584.56 us、N=20 = 14.71 ms,就是這刀切完以後的數字。Dao 在這裡已經不支援,所以沒有外部對手。我們能確定的是:584 us 對上大約 134 MB traffic,有效頻寬約 229 GB/s,已經接近「在搬資料」。再往上加 Tensor Core,先問的不該是 MMA 能不能更快,而是還能不能少一個 pass。
如果第一段沒推進到 15,N=16 的 Pass 2 會剩下 6 層,連 warp shuffle 都走不成,還要再付一次 tiled smem 的 bank conflict 和更窄的 tile。那筆 before 我們沒有單獨留成一張乾淨的表;能確定的是切法本身比較重。
剩下超過 5 層,shuffle 不夠用。這時才回到 tiled shared memory。切之前 Pass 2 一次只搬 16 條 column;切完後一次 64 條,load/store 用 float4,pitch pad 4 個 float,需要時用 ld.global.nc。
這一刀沒有 N=15 那種 3.3 倍的單一數字,因為它是 large-N 的底線,不是主戰場。它要解決的是:既然 global memory 避不掉,就不要再用 16-wide 的 tile 跟不連續 access 自己加罰。Day 17 把 kernel 拆過了;今天只記一句——tile 寬度是 bandwidth 問題,不是美觀問題。
end-to-end 也一樣。Tensor Core 那條路如果只量 MMA、不量 transpose,會得到一張很好看、但不是同一題的速度表。昨天表二的 380 us 對 40 us,差的就是把那些邊算進去之後的結果。
N <= 15
一個 block 做完 ← 第一刀,N=15 從 278 us 到 83 us
N > 15
Pass 1:先做完 15 層
Pass 2:
剩下 <= 5 層 → warp shuffle ← 第二刀
剩下 > 5 層 → 64-wide float4 tile ← 第三刀
不是永遠最佳。但它的每一刀都對得上昨天的表:N=15 贏在少搬一次;N=16 慢在避不掉的那一次;Tensor Core 慢在多出來的那幾次。
N=16N=16 的 584 us 看起來很痛,直覺會想:那把剩下那一層塞進 MMA 呢?
昨天表二已經回答了。TF32 在更小的 N 就因為 transpose 慢九倍。N=16 的下限是 memory traffic。再加 layout conversion,只會把 229 GB/s 那條路弄得更髒。
Tensor Core 該出現的地方,還是資料載入後可以重複使用、每個 output 有大量累加、layout 轉換能被 MMA 攤掉的題。GEMM 是那種題。FWT 不是。
這跟前面 AdaptiveGEMM 不矛盾。同一張卡,GEMM 願意付 Ozaki 的固定成本,是因為乘加夠重。FWT 每次只做一次加減,同一筆帳付不起。
Hadamard 看起來像矩陣乘法
→ 真正該算的是 FWT butterfly
→ 小 N:資料留在一個 block
→ 大 N:才付 2-pass 的 global memory
→ 改寫成 Tensor Core GEMM,layout 比加減法還貴
→ 跟 Dao 放在同一張表上,N=15 切對才贏
→ 重切:278 us → 83 us,N=16 仍然撞在搬資料
如果前面 GEMM 系列是在學怎麼把精度補回這張卡真正有吞吐的地方,這六天就是在學另一件事:
先尊重 operator 的資料流,再談硬體加速。
Hadamard 這一段,先收到這句,以及那一張 before / after 為止。
明天開始 會來一個完全不一樣的題目(因為我沒題目了 阿哈哈哈 而且又有點忙 我決定把僅剩的案例拿出來ㄌ 後續 嗯 希望我有時間把原版的文章寫ㄉ好一點 )
(希望我之後有力氣 把這個部分的PR 以及案例成果寫成好的一篇blog 如果成功 你就會在下面看到我的文章鏈結 26/09/09 隨筆 就交給未來ㄉ我ㄌ)